from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
import torch
import numpy as np
from io import BytesIO
import asyncio
from typing import Dict, Optional
from pydantic import BaseModel

from config import SERVER_HOST, SERVER_PORT, DEVICE
from models.time_series import TimeSeriesLSTM, TimeSeriesTrainer
from models.image_classifier import ImageClassifier, ImageClassifierTrainer
from fl_server import start_server
from fl_client import FederatedClient
from task_manager import TaskManager
from datasets import prepare_etth_data, prepare_helmet_data
from utils.gpu_monitor import GPUMonitor
import threading

app = FastAPI()
task_manager = TaskManager()
gpu_monitor = GPUMonitor()

# CORS中间件配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class TaskCreate(BaseModel):
    task_type: str
    params: Dict

class DatasetUpload(BaseModel):
    dataset_type: str
    dataset_name: str

@app.post("/tasks/create")
async def create_task(task_data: TaskCreate):
    """创建新任务"""
    try:
        task_id = task_manager.create_task(task_data.task_type, task_data.params)
        return {"task_id": task_id, "status": "created"}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@app.get("/tasks/{task_id}")
async def get_task_info(task_id: str):
    """获取任务信息"""
    try:
        return task_manager.get_task_info(task_id)
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))

@app.get("/tasks/{task_id}/log")
async def get_task_log(task_id: str):
    """获取任务日志"""
    try:
        return {"log": task_manager.get_task_log(task_id)}
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))

@app.post("/tasks/{task_id}/start")
async def start_task(task_id: str):
    """启动任务"""
    try:
        task_info = task_manager.get_task_info(task_id)
        task_type = task_info["type"]
        
        if task_type == "time_series":
            asyncio.create_task(run_time_series_task(task_id))
        elif task_type == "image_classification":
            asyncio.create_task(run_image_classification_task(task_id))
        else:
            raise ValueError(f"Unknown task type: {task_type}")
        
        return {"message": f"Task {task_id} started"}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@app.post("/tasks/upload_dataset")
async def upload_dataset(
    file: UploadFile = File(...),
    dataset_type: str = None,
    dataset_name: str = None
):
    """上传数据集"""
    try:
        content = await file.read()
        result = task_manager.save_dataset(content, dataset_type, dataset_name or file.filename)
        return result
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@app.get("/gpu/status")
async def get_gpu_status():
    """获取GPU状态"""
    return gpu_monitor.get_gpu_status()

